#pragma once

#include <initializer_list>
#include <iomanip>
#include <iostream>
#include <lapacke.h>

template <int Row, int Col>
class lapackSolver
{
public:
    // 非常数数组
    lapackSolver(double _A[])
    {
        for (int i = 0; i < Row * Col; A[i] = _A[i], i++)
            ;
    }
    // 用于常数数组
    lapackSolver(std::initializer_list<double> l)
    {
        auto it = l.begin();
        for (int i = 0; i < Row * Col; A[i] = *it++, i++)
            ;
    }

    lapack_int solve(double b[])
    {
        // 标准求解线性方程组
        lapack_int info;
        lapack_int ipiv[Col];

        info = LAPACKE_dgesv(LAPACK_ROW_MAJOR, Col, 1, A, Col, ipiv, b, 1);
        print(Col, b);
        return info;
    }

    lapack_int least_square(double b[])
    {
        // 满秩求解最小二乘问题，使用 QR 分解
        lapack_int info;

        // 得到的结果存放在 b 中 Col x 1
        info = LAPACKE_dgels(LAPACK_ROW_MAJOR, 'N', Row, Col, 1, A, Row, b, 1);
        print(Col, b);
        return info;
    }

    lapack_int eigen(double wr[], double wi[])
    {
        // 求解实矩阵的特征值和特征向量
        LAPACK_D_SELECT2 select;
        lapack_int info, sdim;
        double vs[Row * Row];
        // A 为 Schur 标准型， wr,wi 分别存放特征值的实部和虚部
        info = LAPACKE_dgees(LAPACK_ROW_MAJOR, 'N', 'N', select, Row, A, Row, &sdim, wr, wi, vs, Row);

        std::cout << "eigenvalues: " << std::endl;
        for (int i = 0; i < Row; i++)
        {
            std::cout << wr[i] << " + i " << wi[i] << ", ";
        }
        std::cout << std::endl
                  << "Schur: " << std::endl;
        print(Row, Col, A);
        return info;
    }

    lapack_int eigen_sy(double ev[])
    {
        // 求解实对称矩阵的特征值和特征向量
        lapack_int info;
        // A 存放特征向量， ev 存放特征值
        info = LAPACKE_dsyev(LAPACK_ROW_MAJOR, 'V', 'U', Row, A, Row, ev);
        std::cout << "eigenvalues: " << std::endl;
        print(Row, ev);
        std::cout << "eigenvectors: " << std::endl;
        print(Row, Col, A);
        return info;
    }

    lapack_int Hessenberg()
    {
        // 存放系数
        double tau[Row - 1];
        lapack_int info;
        // 计算矩阵的上 Hessenberg 化，返回 A 包含上 Hessenberg 部分，和下面的变换矩阵部分
        info = LAPACKE_dgehrd(LAPACK_ROW_MAJOR, Row, 1, Row, A, Row, tau);
        std::cout << "Hessnberg: " << std::endl;
        print(Row, Col, A);
        return info;
    }

    lapack_int QR()
    {
        // 存放系数
        double tau[Col];
        lapack_int info;
        // 计算矩阵的 QR 分解
        info = LAPACKE_dgeqrf(LAPACK_ROW_MAJOR, Row, Col, A, Row, tau);
        std::cout << "QR: " << std::endl;
        print(Row, Col, A);
        return info;
    }

    lapack_int inv()
    {
        // 求矩阵的逆
        lapack_int info, ipiv[Row];
        info = LAPACKE_dgetri(LAPACK_ROW_MAJOR, Row, A, Row, ipiv);
        std::cout << "inv: " << std::endl;
        print(Row, Col, A);
        return info;
    }

    // 输出向量和矩阵
    static void print(int dim, double vec[], int w = 0)
    {
        for (int i = 0; i < dim; i++)
        {
            std::cout << std::setw(w) << vec[i] << " ";
        }
        std::cout << std::endl;
    }

    static void print(int row, int col, double mat[], int w = 12)
    {
        for (int i = 0; i < row; i++)
        {
            print(col, &mat[i], w);
        }
        std::cout << std::endl;
    }

    int row() const { return Row; }
    int col() const { return Col; }

private:
    double A[Row * Col];

    // 禁用拷贝构造
    lapackSolver() {}
    lapackSolver(const lapackSolver &) {}
};
